// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <optional>
#include <pollux/testing/gtest_utils.h>
#include <pollux/functions/prestosql/registration/registration_functions.h>
#include <pollux/testing/functions/prestosql/function_base_test.h>
#include <pollux/functions/prestosql/types/timestamp_with_time_zone_type.h>

using namespace kumo::pollux;
using namespace kumo::pollux::test;
using namespace kumo::pollux::exec;
using namespace kumo::pollux::functions::test;

namespace {

class ArrayContainsTest : public FunctionBaseTest {
 public:
  void testContainsConstantKey(
      const ArrayVectorPtr& array_vector,
      const std::vector<int64_t>& search,
      const std::vector<std::optional<bool>>& expected) {
    auto constSearch =
        BaseVector::wrap_in_constant(1, 0, make_array_vector<int64_t>({search}));
    auto result =
        evaluate("contains(c0, c1)", make_row_vector({array_vector, constSearch}));
    assertEqualVectors(make_nullable_flat_vector<bool>(expected), result);
  }

  template <typename T>
  void testContains(
      const ArrayVectorPtr& array_vector,
      T search,
      const std::vector<std::optional<bool>>& expected) {
    auto result = evaluate(
        "contains(c0, c1)",
        make_row_vector({
            array_vector,
            make_constant(search, array_vector->size()),
        }));

    assertEqualVectors(make_nullable_flat_vector<bool>(expected), result);
  };

  void testContainsGeneric(
      const VectorPtr& array_vector,
      const VectorPtr& search,
      const std::vector<std::optional<bool>>& expected) {
    auto result = evaluate(
        "contains(c0, c1)",
        make_row_vector({
            array_vector,
            search,
        }));

    assertEqualVectors(make_nullable_flat_vector<bool>(expected), result);
  };

  template <typename T>
  void testFloatingPointNaNs() {
    static const T kNaN = std::numeric_limits<T>::quiet_NaN();
    static const T kSNaN = std::numeric_limits<T>::signaling_NaN();
    {
      auto array_vector =
          make_array_vector<T>({{1, 2, 3, 4}, {3, 4, kNaN}, {5, 6, 7, 8, kSNaN}});
      // Test fast path for flat.
      testContains(array_vector, kNaN, {false, true, true});
      // Test code path for generic encoded vectors.
      auto indices = make_indices(array_vector->size(), melon::identity);
      auto dictOverArray = wrap_in_dictionary(indices, array_vector);
      testContainsGeneric(
          dictOverArray,
          make_constant(kNaN, array_vector->size()),
          {false, true, true});
    }

    // Test code path for complex-type elements.
    {
      RowTypePtr rowType;
      if constexpr (std::is_same_v<T, float>) {
        rowType = ROW({REAL(), VARCHAR()});
      } else {
        static_assert(std::is_same_v<T, double>);
        rowType = ROW({DOUBLE(), VARCHAR()});
      }
      using ArrayOfRow = std::vector<std::optional<std::tuple<T, std::string>>>;
      std::vector<ArrayOfRow> data = {
          {{{1, "red"}}, {{2, "blue"}}, {{3, "green"}}},
          {{{1, "red"}}, {{kNaN, "blue"}}, {{3, "green"}}},
          {{{1, "red"}}, {{kSNaN, "blue"}}, {{3, "green"}}}};
      auto array_vector = make_array_of_row_vector(data, rowType);
      const auto searchVector =
          make_constant_row(rowType, variant::row({kNaN, "blue"}), 2);
      testContainsGeneric(array_vector, searchVector, {false, true, true});
    }
  }
};

TEST_F(ArrayContainsTest, integerNoNulls) {
  auto array_vector = make_array_vector<int64_t>(
      {{1, 2, 3, 4}, {3, 4, 5}, {}, {5, 6, 7, 8, 9}, {7}, {10, 9, 8, 7}});

  testContains(array_vector, 1, {true, false, false, false, false, false});
  testContains(array_vector, 3, {true, true, false, false, false, false});
  testContains(array_vector, 5, {false, true, false, true, false, false});
  testContains(array_vector, 7, {false, false, false, true, true, true});
  testContains(array_vector, -2, {false, false, false, false, false, false});
  testContains(
      array_vector,
      std::optional<int64_t>(std::nullopt),
      {std::nullopt,
       std::nullopt,
       std::nullopt,
       std::nullopt,
       std::nullopt,
       std::nullopt});
}

TEST_F(ArrayContainsTest, integerWithNulls) {
  auto array_vector = make_nullable_array_vector<int64_t>(
      {{1, 2, 3, 4},
       {3, 4, 5},
       {},
       {5, 6, std::nullopt, 7, 8, 9},
       {7, std::nullopt},
       {10, 9, 8, 7}});

  testContains(
      array_vector, 1, {true, false, false, std::nullopt, std::nullopt, false});
  testContains(
      array_vector, 3, {true, true, false, std::nullopt, std::nullopt, false});
  testContains(array_vector, 5, {false, true, false, true, std::nullopt, false});
  testContains(array_vector, 7, {false, false, false, true, true, true});
  testContains(
      array_vector,
      -2,
      {false, false, false, std::nullopt, std::nullopt, false});
  testContains(
      array_vector,
      std::optional<int64_t>(std::nullopt),
      {std::nullopt,
       std::nullopt,
       std::nullopt,
       std::nullopt,
       std::nullopt,
       std::nullopt});
}

TEST_F(ArrayContainsTest, varcharNoNulls) {
  auto array_vector = make_array_vector<StringView>({
      {"red"_sv, "blue"_sv},
      {"blue"_sv, "yellow"_sv, "orange"_sv},
      {},
      {"red"_sv, "purple"_sv, "green"_sv},
  });

  testContains(array_vector, "red"_sv, {true, false, false, true});
  testContains(array_vector, "blue"_sv, {true, true, false, false});
  testContains(array_vector, "yellow"_sv, {false, true, false, false});
  testContains(array_vector, "green"_sv, {false, false, false, true});
  testContains(array_vector, "crimson red"_sv, {false, false, false, false});
  testContains(
      array_vector,
      std::optional<StringView>(std::nullopt),
      {std::nullopt, std::nullopt, std::nullopt, std::nullopt});
}

TEST_F(ArrayContainsTest, varcharWithNulls) {
  auto array_vector = make_nullable_array_vector<StringView>({
      {"red"_sv, "blue"_sv},
      {std::nullopt, "blue"_sv, "yellow"_sv, "orange"_sv},
      {},
      {"red"_sv, "purple"_sv, "green"_sv},
  });

  testContains(array_vector, "red"_sv, {true, std::nullopt, false, true});
  testContains(array_vector, "blue"_sv, {true, true, false, false});
  testContains(array_vector, "yellow"_sv, {false, true, false, false});
  testContains(array_vector, "green"_sv, {false, std::nullopt, false, true});
  testContains(
      array_vector, "crimson red"_sv, {false, std::nullopt, false, false});
  testContains(
      array_vector,
      std::optional<StringView>(std::nullopt),
      {std::nullopt, std::nullopt, std::nullopt, std::nullopt});
}

TEST_F(ArrayContainsTest, booleanNoNulls) {
  auto array_vector = make_array_vector<bool>({
      {true, false},
      {true},
      {false},
      {},
      {true, false, true},
      {false, false, false},
  });

  testContains(array_vector, true, {true, true, false, false, true, false});
  testContains(array_vector, false, {true, false, true, false, true, true});
  testContains(
      array_vector,
      std::optional<bool>(std::nullopt),
      {std::nullopt,
       std::nullopt,
       std::nullopt,
       std::nullopt,
       std::nullopt,
       std::nullopt});
}

TEST_F(ArrayContainsTest, booleanWithNulls) {
  auto array_vector = make_nullable_array_vector<bool>({
      {true, false},
      {true},
      {false, std::nullopt},
      {},
      {true, false, std::nullopt, true},
      {false, false, false},
  });

  testContains(
      array_vector, true, {true, true, std::nullopt, false, true, false});
  testContains(array_vector, false, {true, false, true, false, true, true});
  testContains(
      array_vector,
      std::optional<bool>(std::nullopt),
      {std::nullopt,
       std::nullopt,
       std::nullopt,
       std::nullopt,
       std::nullopt,
       std::nullopt});
}

TEST_F(ArrayContainsTest, row) {
  std::vector<std::vector<std::optional<std::tuple<int32_t, std::string>>>>
      data = {
          {{{1, "red"}}, {{2, "blue"}}, {{3, "green"}}},
          {{{2, "blue"}}, std::nullopt, {{5, "green"}}},
          {},
          {{{1, "yellow"}}, {{2, "blue"}}, {{4, "green"}}, {{5, "purple"}}},
      };

  auto rowType = ROW({INTEGER(), VARCHAR()});
  auto array_vector = make_array_of_row_vector(data, rowType);

  auto testContains = [&](int32_t n,
                          const char* color,
                          const std::vector<std::optional<bool>>& expected) {
    auto search =
        make_constant_row(rowType, variant::row({n, color}), array_vector->size());

    auto result = evaluate<SimpleVector<bool>>(
        "contains(c0, c1)", make_row_vector({array_vector, search}));

    assertEqualVectors(make_nullable_flat_vector<bool>(expected), result);
  };

  testContains(1, "red", {true, std::nullopt, false, false});
  testContains(2, "blue", {true, true, false, true});
  testContains(4, "green", {false, std::nullopt, false, true});
  testContains(5, "green", {false, true, false, false});
  testContains(1, "purple", {false, std::nullopt, false, false});
}

TEST_F(ArrayContainsTest, preDefinedResults) {
  auto array_vector = make_array_vector<int64_t>(
      {{1, 2, 3, 4}, {3, 4, 5}, {}, {5, 6, 7, 8, 9}, {7}, {10, 9, 8, 7}});

  testContains(array_vector, 1, {true, false, false, false, false, false});
  testContains(array_vector, 1, {true, false, false, false, false, false});
  testContains(array_vector, 3, {true, true, false, false, false, false});
  testContains(array_vector, 5, {false, true, false, true, false, false});
  testContains(array_vector, 7, {false, false, false, true, true, true});
  testContains(array_vector, -2, {false, false, false, false, false, false});
  testContains(
      array_vector,
      std::optional<int64_t>(std::nullopt),
      {std::nullopt,
       std::nullopt,
       std::nullopt,
       std::nullopt,
       std::nullopt,
       std::nullopt});
}

TEST_F(ArrayContainsTest, preAllocatedNulls) {
  auto array_vector = make_array_vector<int64_t>(
      {{1, 2, 3, 4}, {3, 4, 5}, {}, {5, 6, 7, 8, 9}, {7}, {10, 9, 8, 7}});

  auto testContains = [&](std::optional<int64_t> search,
                          const std::vector<std::optional<bool>>& expected) {
    VectorPtr result = make_flat_vector<bool>(6);
    SelectivityVector rows(6);
    rows.resize(6);
    result->set_null(0, true);

    evaluate<SimpleVector<bool>>(
        "contains(c0, c1)",
        make_row_vector({
            array_vector,
            make_constant(search, array_vector->size()),
        }),
        rows,
        result);

    assertEqualVectors(make_nullable_flat_vector<bool>(expected), result);
  };

  testContains(1, {true, false, false, false, false, false});
  testContains(3, {true, true, false, false, false, false});
  testContains(5, {false, true, false, true, false, false});
  testContains(7, {false, false, false, true, true, true});
  testContains(-2, {false, false, false, false, false, false});
  testContains(
      std::nullopt,
      {std::nullopt,
       std::nullopt,
       std::nullopt,
       std::nullopt,
       std::nullopt,
       std::nullopt});
}

TEST_F(ArrayContainsTest, constantEncodingElements) {
  // ArrayVector with ConstantVector<Array> elements.
  auto baseVector = make_array_vector<int64_t>(
      {{1, 2, 3, 4}, {3, 4, 5}, {6}, {5, 6, 7, 8, 9}, {7}, {10, 9, 8, 7}});
  const vector_size_t kTopLevelVectorSize = baseVector->size() * 2;
  auto constant_vector =
      BaseVector::wrap_in_constant(kTopLevelVectorSize, 0, baseVector);
  auto array_vector = make_array_vector({0, 3}, constant_vector);

  testContainsConstantKey(array_vector, {1, 2, 3, 4}, {true, true});
  testContainsConstantKey(array_vector, {3, 4}, {false, false});
  testContainsConstantKey(array_vector, {5, 6, 7, 8, 9}, {false, false});
}

TEST_F(ArrayContainsTest, dictionaryEncodingElements) {
  // ArrayVector with DictionaryVector<Array> elements.
  auto baseVector =
      make_array_vector<int64_t>({{1, 2, 3, 4}, {3, 4, 5}, {10, 9, 8, 7}});
  auto baseVectorSize = baseVector->size();
  const vector_size_t kTopLevelVectorSize = baseVectorSize * 2;
  BufferPtr indices = allocate_indices(kTopLevelVectorSize, pool_.get());
  auto rawIndices = indices->asMutable<vector_size_t>();
  for (size_t i = 0; i < kTopLevelVectorSize; ++i) {
    rawIndices[i] = i % baseVectorSize;
  }
  auto dictVector = BaseVector::wrap_in_dictionary(
      nullptr, indices, kTopLevelVectorSize, baseVector);
  auto array_vector = make_array_vector({0, baseVectorSize + 1}, dictVector);
  // array_vector is
  // {
  //    [[1, 2, 3, 4], [3, 4, 5], [10, 9, 8, 7], [1, 2, 3, 4]],
  //    [[3, 4, 5], [10, 9, 8, 7]]
  // }
  testContainsConstantKey(array_vector, {1, 2, 3, 4}, {true, false});
  testContainsConstantKey(array_vector, {3, 4, 5}, {true, true});
}

TEST_F(ArrayContainsTest, arrayCheckNulls) {
  kumo::pollux::functions::prestosql::registerInternalFunctions();

  static const std::string kErrorMessage =
      "contains does not support arrays with elements that contain null";
  auto contains = [&](const std::string& search,
                      const auto& data,
                      bool internal = false) {
    const auto searchBase = make_array_vector_from_json<int32_t>({search});
    const auto searchConstant =
        BaseVector::wrap_in_constant(data->size(), 0, searchBase);
    std::string call = internal ? "\"$internal$contains\"" : "contains";
    const auto result = evaluate(
        fmt::format("{}(c0, c1)", call), make_row_vector({data, searchConstant}));
    return result->template as_flat_vector<bool>()->value_at(0);
  };

  {
    // Null at the end of the array.
    const auto baseVector = make_array_vector_from_json<int32_t>({
        "[1, 1]",
        "[2, 2]",
        "[3, null]",
        "[4, 4]",
        "[5, 5]",
        "[6, 6]",
    });
    const auto data = make_array_vector({0, 3}, baseVector);

    // No null equal.
    ASSERT_FALSE(contains("[7, null]", data));
    ASSERT_FALSE(contains("[7, null]", data, true));
    // Null equal, [3, null] vs [3, 3].
    POLLUX_ASSERT_THROW(contains("[3, 3]", data), kErrorMessage);
    ASSERT_FALSE(contains("[3, 3]", data, true));
    // Null equal, [6, 6] vs [6, null].
    POLLUX_ASSERT_THROW(contains("[6, null]", data), kErrorMessage);
    ASSERT_FALSE(contains("[6, null]", data, true));
    // [3, null] = [3, null] is true in $internal$contains.
    ASSERT_TRUE(contains("[3, null]", data, true));
  }

  {
    // Null at the beginning of the array.
    // data is [[null, 3]].
    const auto data =
        make_nullable_nested_array_vector<int32_t>({{{{{std::nullopt, 3}}}}});

    // [null] = [null, 3] is false.
    ASSERT_FALSE(contains("[null]", data));
    ASSERT_FALSE(contains("[null]", data, true));
    //  [null, 4] = [null, 3] is false.
    ASSERT_FALSE(contains("[null, 4]", data));
    ASSERT_FALSE(contains("[null, 4]", data, true));
    //  [null, 4] = [1, 1] is false.
    ASSERT_FALSE(contains("[1, 1]", data));
    ASSERT_FALSE(contains("[1, 1]", data, true));

    // [null, 3] = [null, 3] is indeterminate.
    POLLUX_ASSERT_THROW(contains("[null, 3]", data), kErrorMessage);
    // [null, 3] = [null, null] is indeterminate.
    POLLUX_ASSERT_THROW(contains("[null, null]", data), kErrorMessage);
    ASSERT_FALSE(contains("[null, null]", data, true));
    // [null, 3] = [null, 3] is true in $internal$contains.
    ASSERT_TRUE(contains("[null, 3]", data, true));
  }
}

TEST_F(ArrayContainsTest, rowCheckNulls) {
  kumo::pollux::functions::prestosql::registerInternalFunctions();

  const auto baseVector = make_row_vector({
      make_nullable_flat_vector<int32_t>({1, 2, 3, 4, 5, 6}),
      make_nullable_flat_vector<int32_t>({1, 2, std::nullopt, 4, 5, 6}),
  });
  const auto data = make_array_vector({0, 3}, baseVector);

  auto contains = [&](const std::vector<std::optional<int32_t>>& search,
                      bool internal = false) {
    const auto searchBase = make_row_vector({
        make_nullable_flat_vector<int32_t>({search.at(0)}),
        make_nullable_flat_vector<int32_t>({search.at(1)}),
    });
    const auto searchConstant =
        BaseVector::wrap_in_constant(data->size(), 0, searchBase);
    std::string call = internal ? "\"$internal$contains\"" : "contains";
    const auto result = evaluate(
        fmt::format("{}(c0, c1)", call), make_row_vector({data, searchConstant}));
    return result->as_flat_vector<bool>()->value_at(0);
  };

  static const std::string kErrorMessage =
      "contains does not support arrays with elements that contain null";
  // No null equal.
  ASSERT_FALSE(contains({7, std::nullopt}));
  ASSERT_FALSE(contains({7, std::nullopt}, true));
  // Null equal, (3, null) vs (3, 3).
  POLLUX_ASSERT_THROW(contains({3, 3}), kErrorMessage);
  ASSERT_FALSE(contains({3, 3}, true));
  // Null equal, (6, 6) vs (6, null).
  POLLUX_ASSERT_THROW(contains({6, std::nullopt}), kErrorMessage);
  ASSERT_FALSE(contains({6, std::nullopt}, true));
  // (3, null) = (3, null) is true in $internal$contains.
  ASSERT_TRUE(contains({3, std::nullopt}, true));
}

TEST_F(ArrayContainsTest, floatNaNs) {
  testFloatingPointNaNs<float>();
  testFloatingPointNaNs<double>();
}

TEST_F(ArrayContainsTest, timestampWithTimeZone) {
  auto array_vector = make_array_vector(
      {0, 4, 7, 7, 13, 15},
      make_nullable_flat_vector<int64_t>(
          {pack(1, 1),
           pack(2, 2),
           pack(3, 3),
           pack(4, 4),
           pack(3, 5),
           pack(4, 6),
           pack(5, 7),
           pack(5, 8),
           pack(6, 9),
           std::nullopt,
           pack(7, 10),
           pack(8, 11),
           pack(9, 12),
           pack(7, 13),
           std::nullopt,
           pack(10, 14),
           pack(9, 15),
           pack(8, 16),
           pack(7, 17)},
          TIMESTAMP_WITH_TIME_ZONE()));

  const auto testContains =
      [&](std::optional<int64_t> needle,
          const std::vector<std::optional<bool>>& expected) {
        const auto searchVector = make_constant(
            needle, array_vector->size(), TIMESTAMP_WITH_TIME_ZONE());

        testContainsGeneric(array_vector, searchVector, expected);
      };

  testContains(
      pack(1, 1), {true, false, false, std::nullopt, std::nullopt, false});
  testContains(
      pack(3, 3), {true, true, false, std::nullopt, std::nullopt, false});
  testContains(pack(5, 1), {false, true, false, true, std::nullopt, false});
  testContains(pack(7, 2), {false, false, false, true, true, true});
  testContains(
      pack(-2, 1), {false, false, false, std::nullopt, std::nullopt, false});
  testContains(
      std::nullopt,
      {std::nullopt,
       std::nullopt,
       std::nullopt,
       std::nullopt,
       std::nullopt,
       std::nullopt});

  // Test wrapped in a complex value.
  array_vector = make_array_vector(
      {0, 4, 7, 7, 12, 13},
      make_row_vector({make_nullable_flat_vector<int64_t>(
          {pack(1, 1),
           pack(2, 2),
           pack(3, 3),
           pack(4, 4),
           pack(3, 5),
           pack(4, 6),
           pack(5, 7),
           pack(5, 8),
           pack(6, 9),
           pack(7, 10),
           pack(8, 11),
           pack(9, 12),
           pack(7, 13),
           pack(10, 14),
           pack(9, 15),
           pack(8, 16),
           pack(7, 17)},
          TIMESTAMP_WITH_TIME_ZONE())}));

  const auto testContainsRow =
      [&](int64_t needle, const std::vector<std::optional<bool>>& expected) {
        const auto searchVector = BaseVector::wrap_in_constant(
            array_vector->size(),
            0,
            make_row_vector({make_flat_vector(
                std::vector<int64_t>{needle}, TIMESTAMP_WITH_TIME_ZONE())}));

        testContainsGeneric(array_vector, searchVector, expected);
      };

  testContainsRow(pack(1, 1), {true, false, false, false, false, false});
  testContainsRow(pack(3, 3), {true, true, false, false, false, false});
  testContainsRow(pack(5, 1), {false, true, false, true, false, false});
  testContainsRow(pack(7, 2), {false, false, false, true, true, true});
  testContainsRow(pack(-2, 1), {false, false, false, false, false, false});
}
} // namespace
